#!/usr/bin/env python2.5

## Trains an HMM tagger upon our data.

## Usage:  ~ sentences num.sents

import sys, nltk
from nltk import tag
from nltk.tag import hmm
from nltk.tag.hmm import *

import tokenize

class getHMMTagger:
	"""Creates, trains, and returns a Hidden Markov Model tagger."""

	def __init__(self, file, NUM_SENTS):
		"""Create the tagger."""
		tagfile = open(file, 'r')
		self.seq = list()
		sent = list()
		self.token_seq = list()
		self.tag_seq = list()
		RCE = 0
		for entry in tagfile:
			if ((NUM_SENTS != 0) and (NUM_SENTS == RCE)):	break
			if (entry == "\n"):
				self.seq.append(sent)
				sent = list()
				continue
			entry = tokenize.tokenize(entry, 3)
			sent.append([" ".join(entry.split()[:-1]), entry.split()[-1]])
			self.token_seq.append(" ".join(entry.split()[:-1]))
			self.tag_seq.append(entry.split()[-1])
			RCE += 1
		self.HMMTagger = HiddenMarkovModelTrainer(self.tag_seq, self.token_seq)

	def getSentences(self):
		"""Prints the sentences which we used to create and train the tagger."""
		for sent in self.seq:	print sent

	def trainHMM(self):
		"""Train the tagger on the sequence of sentences from the input file."""
		return self.HMMTagger.train_supervised(self.seq)

## Run a simple test.
if (__name__ == "__main__"):
	## Run this from a class definition, for testing purposes
	HMMT = getHMMTagger(sys.argv[1], int(sys.argv[2]))
	# HMMT.getSentences()
	HMMTrained = HMMT.trainHMM()
	# test = [["This", "needs", "more", "crackers", "."]]
	test = list(["Inertia is a property of the Mars 's matter .".split(" ")])
	print HMMTrained.batch_tag(test)
	test = list(["This is a brown fungi in the summer .".split(" ")])
	print HMMTrained.batch_tag(test)
